package com.wfh.aiagent.controller;

import com.wfh.aiagent.app.ProgramingApp;
import com.wfh.aiagent.cache.RedisSemanticCacheService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Optional;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @Author FengHuan Wang
 * @Date 2025/5/25 14:51
 * @Version 1.0
 */
@Slf4j
@RestController
@RequestMapping("/app")
public class AppController {

    @Resource
    private ProgramingApp programingApp;

    @Resource
    private RedisSemanticCacheService redisSemanticCacheService;

    public record ChatRequest(String message, String chatId) {
    }

    @PostMapping(value = "/chat")
    public SseEmitter chat(@RequestBody ChatRequest chatRequest) {
        SseEmitter emitter = new SseEmitter(120_000L); // 设置2分钟超时

        log.info("收到聊天请求: chatId={}, message={}", chatRequest.chatId, chatRequest.message);

        // 1. 先尝试从缓存获取
        Optional<String> cachedAnswer = redisSemanticCacheService.get(chatRequest.message);
        if (cachedAnswer.isPresent()) {
            log.info("缓存命中，直接返回缓存内容");
            try {
                // 发送缓存的内容
                emitter.send(SseEmitter.event()
                        .data(cachedAnswer.get())
                        .name("message"));
                emitter.complete();
            } catch (IOException e) {
                log.error("发送缓存内容失败", e);
                emitter.completeWithError(e);
            }
            return emitter;
        }

        log.info("缓存未命中，开始流式生成响应");

        // 使用线程安全的集合来收集响应内容
        CopyOnWriteArrayList<String> responseChunks = new CopyOnWriteArrayList<>();
        AtomicReference<String> fullResponse = new AtomicReference<>("");

        try {
            // 2. 流式生成响应
            programingApp.doChatWithRagToolStream(chatRequest.message, chatRequest.chatId)
                    .subscribe(
                            chunk -> {
                                try {
                                    // 收集每个chunk
                                    responseChunks.add(chunk);
                                    fullResponse.updateAndGet(current -> current + chunk);

                                    // 实时发送给客户端
                                    emitter.send(SseEmitter.event()
                                            .data(chunk)
                                            .name("message"));
                                    log.debug("发送chunk: {}", chunk);
                                } catch (IOException e) {
                                    log.error("发送chunk失败", e);
                                    throw new RuntimeException("SSE发送失败", e);
                                }
                            },
                            error -> {
                                log.error("流式响应生成失败", error);
                                try {
                                    emitter.send(SseEmitter.event()
                                            .data("抱歉，生成回答时出现错误")
                                            .name("error"));
                                    emitter.completeWithError(error);
                                } catch (IOException e) {
                                    log.error("发送错误消息失败", e);
                                }
                            },
                            () -> {
                                try {
                                    // 流式响应完成
                                    String completeResponse = fullResponse.get();
                                    log.info("流式响应完成，总长度: {}", completeResponse.length());

                                    // 发送完成事件
                                    emitter.send(SseEmitter.event()
                                            .data("[DONE]")
                                            .name("complete"));

                                    // 缓存完整响应
                                    if (!completeResponse.trim().isEmpty()) {
                                        try {
                                            redisSemanticCacheService.put(chatRequest.message, completeResponse);
                                            log.debug("响应已缓存");
                                        } catch (Exception cacheException) {
                                            log.warn("缓存响应失败，但不影响主流程", cacheException);
                                        }
                                    }

                                    emitter.complete();
                                    log.info("SSE连接正常完成");
                                } catch (IOException e) {
                                    log.error("完成SSE连接时发生错误", e);
                                }
                            }
                    );

            // 3. 设置SSE连接的超时和错误处理
            emitter.onTimeout(() -> {
                log.warn("SSE连接超时, chatId: {}", chatRequest.chatId);
                // 即使超时，也尝试缓存已收集的内容
                cachePartialResponse(chatRequest.message, fullResponse.get());
            });

            emitter.onError(error -> {
                log.error("SSE连接发生错误, chatId: {}", chatRequest.chatId, error);
                // 发生错误时也尝试缓存已收集的内容
                cachePartialResponse(chatRequest.message, fullResponse.get());
            });

            emitter.onCompletion(() -> {
                log.info("SSE连接完成, chatId: {}", chatRequest.chatId);
            });

        } catch (Exception e) {
            log.error("处理聊天请求时发生异常", e);
            try {
                emitter.send(SseEmitter.event()
                        .data("系统繁忙，请稍后重试")
                        .name("error"));
                emitter.completeWithError(e);
            } catch (IOException ex) {
                log.error("发送错误消息失败", ex);
            }
        }

        return emitter;
    }

    /**
     * 缓存部分响应内容
     */
    private void cachePartialResponse(String question, String response) {
        if (response != null && !response.trim().isEmpty()) {
            try {
                redisSemanticCacheService.put(question, response);
                log.debug("已缓存部分响应内容，长度: {}", response.length());
            } catch (Exception e) {
                log.warn("缓存部分响应失败", e);
            }
        }
    }

    @GetMapping("/health")
    public String health() {
        return "ok";
    }

    /**
     * 获取缓存状态信息
     */
    @GetMapping("/cache/status")
    public Object getCacheStatus() {
        return redisSemanticCacheService.getStats();
    }
}